package dev.alm.langchain4jspringbootdemo.config;

import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class RagConfig {

    @Value("${langchain4j.community.dashscope.chat-model.api-key}")
    private String apiKey;

    public interface RagDemoService {
        String chat(@MemoryId String memoryId, @UserMessage String message);

        TokenStream stream(@MemoryId String memoryId, @UserMessage String message);
    }

    @Bean
    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean
    public RagDemoService ragDemoService(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel streamingQwenChatModel, InMemoryEmbeddingStore embeddingStore, QwenEmbeddingModel embeddingModel) {
        Document document = ClassPathDocumentLoader.loadDocument("info.txt");
        DocumentBySentenceSplitter documentBySentenceSplitter = new DocumentBySentenceSplitter(50, 10);
        List<TextSegment> segments = documentBySentenceSplitter.split(document);
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        embeddingStore.addAll(embeddings, segments);
        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder().embeddingModel(embeddingModel).embeddingStore(embeddingStore).maxResults(1).minScore(0.7).build();
        RagDemoService service = AiServices.builder(RagDemoService.class).chatLanguageModel(qwenChatModel).streamingChatLanguageModel(streamingQwenChatModel).contentRetriever(retriever).chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).id(memoryId).build()).build();
        return service;
    }

}
